package formal.tree.recursion;

import java.util.HashMap;
import java.util.Map;

/**
 * @author DengYuan2
 * @create 2021-01-13 16:14
 */
public class M_437 {
    public static void main(String[] args) {
        Integer[] a = {10, 5, -3, 3, 2, null, 11, 3, -2, null, 1};
        TreeNode treeNode = TreeNode.generateTree(a);
        int res = pathSum3(treeNode, 7);
        System.out.println(res);
    }

    /**
     * 我的写法-递归-两层-从上到下
     */
    private static int count = 0;

    public static int pathSum(TreeNode root, int sum) {
        if (root == null) {
            return 0;
        }
        helper(root, sum);
        return count;
    }

    public static void helper(TreeNode root, int sum) {
        if (root == null) {
            return;
        }
        sumHelper(root, sum);
        helper(root.left, sum);
        helper(root.right, sum);
    }

    public static void sumHelper(TreeNode root, int sum) {
        if (root == null) {
            return;
        }
        sum -= root.val;
        if (sum == 0) {
            count++;
        }
        sumHelper(root.left, sum);
        sumHelper(root.right, sum);
    }


    /**
     * 大神-将我的方法改进了许多
     *
     * @param root
     * @param sum
     * @return
     */
    public static int pathSum2(TreeNode root, int sum) {
        if (root == null) {
            return 0;
        }
        int res = 0;
        res = pathSumHelper(root, sum) + pathSum2(root.left, sum) + pathSum2(root.right, sum);
        return res;

    }

    public static int pathSumHelper(TreeNode root, int sum) {
        if (root == null) {
            return 0;
        }
        int res = 0;
        sum -= root.val;
        if (sum == 0) {
            res++;
        }
        res = res + pathSumHelper(root.left, sum) + pathSumHelper(root.right, sum);
        return res;
    }

    /**
     * 官方- 前缀和+递归+回溯
     *
     * @param root
     * @param sum
     * @return
     */
    public static int pathSum3(TreeNode root, int sum) {
        HashMap<Integer, Integer> prefixSumCount = new HashMap<>(16);
        prefixSumCount.put(0, 1);
        return recursionPathSum(root,prefixSumCount,sum,0);
    }

    public static int recursionPathSum(TreeNode node, Map<Integer, Integer> prefixSumCount, int target, int curSum) {
        if (node == null) {
            return 0;
        }
        int res = 0;
        curSum += node.val;
        res += prefixSumCount.getOrDefault(curSum - target, 0);
        prefixSumCount.put(curSum, prefixSumCount.getOrDefault(curSum, 0) + 1);
        res += recursionPathSum(node.left, prefixSumCount, target, curSum);
        res += recursionPathSum(node.right, prefixSumCount, target, curSum);

        prefixSumCount.put(curSum, prefixSumCount.get(curSum) - 1);
        return res;
    }
}
